from RumdetecFramework.BaseRumorFramework import RumorDetection
from Dataloader.twitterloader import *
from Dataloader.dataloader_utils import *
from SentModel.Sent2Vec import *
from PropModel.SeqPropagation import GRUModel

def obtain_LSTMRD():
    lvec = W2VLSTMVec("../saved/glove_en/", 300, 1, None, False)
    prop = GRUModel(300, 256, 1, 0.2)
    cls = nn.Linear(256, 2)
    LSTMRD = RumorDetection(lvec, prop, cls)
    return LSTMRD

def obtain_general_set(tr_prefix, dev_prefix, te_prefix):
    tr_set = TwitterSet()
    tr_set.load_data_fast(data_prefix=tr_prefix, min_len=5)
    dev_set = TwitterSet()
    dev_set.load_data_fast(data_prefix=dev_prefix)
    te_set = TwitterSet()
    te_set.load_data_fast(data_prefix=te_prefix)
    return tr_set, dev_set, te_set

i = 0
tr, dev, te = obtain_general_set("../data/twitter_tr%d"%i, "../data/twitter_dev%d"%i, "../data/twitter_te%d"%i)
tr, dev, te = shuffle_data(tr, dev, te)
tr.filter_short_seq(min_len=5)
tr.trim_long_seq(10)
print("%s : (dev event)/(test event)/(train event) = %3d/%3d/%3d" % (te.data[te.data_ID[0]]['event'], len(dev), len(te), len(tr)))
print("\n\n===========SubRDM Train===========\n\n")
model = obtain_LSTMRD()
model.train_iters(tr, dev, te, max_epochs=100,
                log_dir="../logs/", log_suffix="_RDM",
                model_file="W2VRDM_General.pkl")
# model.train_iters(tr, dev, te, max_epochs=100,
#                 log_dir="../logs/", log_suffix="_BertRD",
#                 model_file="W2VLSTMRD_%s.pkl"% (te.data[te.data_ID[0]]['event']))
